import operator
from math import log


#计算熵
def calc_shannon_ent(dataSet):
    numentries = len(dataSet)
    label_counts = {}
    for featVec in dataSet:
        currentLabel = featVec[-1]
        if currentLabel not in label_counts.keys():
            label_counts[currentLabel] = 0
        label_counts[currentLabel] += 1
    shnnonEnt = 0.0
    for key in label_counts:
        prob = float(label_counts[key])/numentries
        shnnonEnt -= prob * log(prob, 2)
    return shnnonEnt


#划分数据集
def split_dataset(dataset, axis, value):
    ret_dataset = []
    for featVec in dataset:
        #print(featVec)
        if featVec[axis] == value:
            reduced_featvec = featVec[:axis]
            reduced_featvec.extend(featVec[axis+1:])
            ret_dataset.append(reduced_featvec)
    return ret_dataset


#选择最好的特征划分
def choose_best_feature_tosplit(dataset):
    num_features = len(dataset[0]) - 1
    base_ent = calc_shannon_ent(dataset)
    best_info_gain = 0.0
    best_feature = -1
    for i in range(num_features):
        list_feature = [example[i] for example in dataset]
        unique_vals = set(list_feature)
        new_ent = 0.0
        for val in unique_vals:
            sub_dataset = split_dataset(dataset, i, val)
            prob = len(sub_dataset)/float(len(dataset))
            new_ent += prob * calc_shannon_ent(sub_dataset)
        info_gain = base_ent - new_ent
        if(info_gain > best_info_gain):
            best_info_gain = info_gain
            best_feature = i
    return best_feature


def major_cnt(list_class):
    class_count = {}
    for vote in list_class:
        if vote not in class_count.keys():
            class_count[vote] = 0
        class_count[vote] += 1
    sorted_class_count = sorted(class_count.items(), key = operator.itemgetter(1), reverse = True)
    return sorted_class_count[0][0]


def create_tree(dataset, labels):
    list_class = [example[-1] for example in dataset]
    if list_class.count(list_class[0]) == len(list_class):
        return list_class
    if len(dataset[0]) == 1:
        return major_cnt(list_class)

    best_feat = choose_best_feature_tosplit(dataset)
    best_feat_label = labels[best_feat]
    mytree = {best_feat_label : {}}
    del(labels[best_feat])
    feat_vals = [example[best_feat] for example in dataset]
    unique_vals = set(feat_vals)
    for val in unique_vals:
        sub_labels = labels[:]
        mytree[best_feat_label][val] = create_tree(split_dataset(dataset, best_feat, val), sub_labels)

    return mytree


#test decision tree
def classify(input_tree, feat_labels, testvec):
    print(input_tree)
    first_side = list(input_tree.keys())
    first_str = first_side[0]
    second_dict = input_tree[first_str]
    feat_index = feat_labels.index(first_str)
    class_label = ''
    for key in second_dict.keys():
        if testvec[feat_index] == key:
            if type(second_dict[key]).__name__ == 'dict':
                class_label = classify(second_dict[key], feat_labels, testvec)
            else:
                class_label = second_dict[key]
    return class_label


#把树存到文件里
def store_tree(input_tree, filename='dtree.txt'):
    import pickle
    fw = open(filename, 'wb')
    pickle.dump(input_tree, fw)
    fw.close()


def grab_tree(filename='dtree.txt'):
    import pickle
    fr = open(filename, 'rb')
    return pickle.load(fr, encoding='utf-8')

if __name__ == '__main__':
    dataSet = [
        [1,1,'yes'],
        [0,1,'no'],
        [1,0,'no'],
        [0,1,'yes']
     ]
    dataSet1 = [
        [1, 1, 'a'],
        [1, 1, 'b'],
        [1, 0, 'c'],
        [0, 1, 'd']
     ]

    test_dataset = [
        [1, 1, 1,'ye'],
        [1, 1, 2,'ye'],
        [1, 1, 3,'no'],
        [0, 1, 1,'no'],
        [0, 0, 2,'no']
    ]

    labels = ['A', 'B', 'C']

    # print(calc_shannon_ent(dataSet1))
    # print(split_dataset(test_dataset, 0, 0))
    #print(choose_best_feature_tosplit(test_dataset))
    #print(create_tree(test_dataset, labels))

    intree = {'A': {0: ['no', 'no'], 1: {'C': {1: ['ye'], 2: ['ye'], 3: ['no']}}}}
    # flabels = ['A', 'C']
    # test = [1, 1]
    # print(list(intree.keys())[0])
    # print(list(flabels).index('A'))
    # print(classify(intree, flabels, test))

    store_tree(intree, 'dtree.txt')
    print(grab_tree())